package cn.edu.zjut;


import lombok.Data;
import org.ujmp.core.Matrix;
import org.ujmp.core.calculation.Calculation;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @ClassName: UKFWindow2
 * @Description:
 * @Author: 12071
 * @Date: 2022/1/23 20:33
 */
@Data
public class FilterWindow2 {
    private String bandMac;
    private List<NodePosition> positionHistory; //定位历史
    private Integer exceptionTime;  // 异常计算次数
    private Long timeSpan; // 采样时间间隔

    private final static int STATE_DIMENSION = 4;     //状态量维数 x,y,vx,vy
    private final static int MEASURE_DIMENSION = 2;     //观测量维度 x,y

    private final static Matrix Q = Matrix.Factory.importFromArray(new double[][]{{2.5*0.000000000001,0},{0,2.5*0.000000000001}});    // 过程噪声方差
    private final static Matrix R = Matrix.Factory.importFromArray(new double[][]{{0.0001*0.0001,0},{0,0.0001*0.0001}});    // 观测噪声方差
    private final static Matrix P_0 = Matrix.Factory.importFromArray(new double[][]{{0.0001*0.0001,0,0,0},{0,0.0001*0.0001,0,0},{0,0,0.0001*0.0001,0},{0,0,0,0.0001*0.0001}});    //协方差矩阵初值

    // UT变换
    private final static double alpha = 1; //default, tunable
    private final static double kalpha = 1; //default, tunable
    private final static double beta = 1.0; //default, tunable
    private static double lambda = alpha * alpha * (STATE_DIMENSION + kalpha) - STATE_DIMENSION;   //数学研究表明最佳值一般为3-N
    private static double weight_0_m = lambda / (STATE_DIMENSION + lambda);
    private static double weight_0_c = lambda / (STATE_DIMENSION + lambda) + (1 - alpha * alpha + beta);
    private static double weight_i = 1 / (2 * (STATE_DIMENSION + lambda));
    private static double ns = Math.sqrt(STATE_DIMENSION + lambda);

    private Matrix B = Matrix.Factory.zeros(STATE_DIMENSION, MEASURE_DIMENSION); // 过程噪声驱动矩阵(动态)
    private Matrix F = Matrix.Factory.zeros(STATE_DIMENSION, STATE_DIMENSION); // 状态转移矩阵（动态）
    private Matrix P = Matrix.Factory.importFromArray(new double[][]{{1000, 0, 0, 0}, {0, 1000, 0, 0}, {0, 0, 0.1 * 2.5, 0}, {0, 0, 0, 0.1 * 2.5}});    //协方差矩阵，初始化
    private Matrix K = Matrix.Factory.zeros(STATE_DIMENSION, MEASURE_DIMENSION);   //Kalman增益系数


    /**
     * 动态调整状态转移矩阵F和过程噪声驱动矩阵B（根据时间差）
     * 建议时间差为固定值，动态存在隐患
     *
     * @param T
     */
    private void calBAndF(long T) {
        this.B = Matrix.Factory.importFromArray(new double[][]{{Math.pow(T, 2) / 2, 0}, {0, Math.pow(T, 2) / 2}, {T, 0}, {0, T}});
        this.F = Matrix.Factory.importFromArray(new double[][]{{1, 0, T, 0}, {0, 1, 0, T}, {0, 0, 1, 0}, {0, 0, 0, 1}});

    }

    public FilterWindow2(String bandMac, Long timeSpan) {
        this.bandMac = bandMac;
        this.positionHistory = new ArrayList<>();
        this.timeSpan = timeSpan;
    }

    /**
     * 获取最新的位置
     *
     * @return
     */
    public NodePosition getLastPosition() {
        int size = positionHistory.size();
        if (size < 4) {
            return null;
        } else {
            return positionHistory.get(size - 1);
        }
    }

    /**
     * 直接添加位置信息
     *
     * @param position
     */
    public void insertPosition(NodePosition position) {
        this.positionHistory.add(position);
    }

    /**
     * 通过滤波添加位置信息
     *
     * @param position
     */
    public void insertPositionWithUKF(NodePosition position) {
        System.out.println("Before UKF:" + position);
        NodePosition fixedPosition = UKFilter(position);
        System.out.println("After UKF:" + fixedPosition);
        insertPosition(fixedPosition);
    }

    /**
     * 根据历史位置估计速度
     *
     * @param time
     * @return
     */
    public double[] estimateSpeed(long time) {
        int timeDelay = 5000 * 3; // 进行计算的时间范围
        List<NodePosition> history = positionHistory.stream().filter(position -> (time - position.getUpdateTime() < timeDelay)).collect(Collectors.toList());
        int size = history.size();
        double vx, vy;
        if (size < 2) {
            vx = 0;
            vy = 0;
        } else {
            vx = (history.get(size - 1).getX() - history.get(0).getX()) / 5;
            vy = (history.get(size - 1).getY() - history.get(0).getY()) / 5;
        }
        return new double[]{vx, vy};
    }

    /**
     * 无迹卡尔曼滤波过程
     */
    public NodePosition UKFilter(NodePosition measurePosition) {
        NodePosition lastPosition = getLastPosition();
        long T = lastPosition != null ? measurePosition.getUpdateTime() - lastPosition.getUpdateTime() : 0;
        if (T == 0 || T > 25000) { //初始数据或过时数据不予判断
            this.P = P_0;
            return measurePosition;
        }
        double[] speed = estimateSpeed(measurePosition.getUpdateTime());
        calBAndF(timeSpan);   // 动态调整转移矩阵
        // 1. 获取Sigma点
        System.out.println("P");
        System.out.println(P);
        // P_0.print(P_0.getRowDimension(), P_0.getColumnDimension());

        //   CholeskyDecomposition chol = P_0.chol();
        Matrix chol = P_0.chol();
        Matrix cholesky = P.chol().times(Calculation.Ret.NEW, true, ns);
        Matrix sigmaPoint = Matrix.Factory.zeros(STATE_DIMENSION, 2 * STATE_DIMENSION + 1);
        Matrix sigma_0 = Matrix.Factory.importFromArray(new double[][]{{lastPosition.getX(), lastPosition.getY(), speed[0], speed[1]}}).transpose();
        for (int i = 0; i < 2 * STATE_DIMENSION + 1; i++) {
            if (i == 0) {
                setColumn(sigmaPoint, i, sigma_0);
            } else if (i < STATE_DIMENSION + 1) {
                setColumn(sigmaPoint, i, sigma_0.plus(cholesky.selectColumns(Calculation.Ret.NEW, i - 1)));
            } else {
                setColumn(sigmaPoint, i, sigma_0.minus(cholesky.selectColumns(Calculation.Ret.NEW, (i - 1) % STATE_DIMENSION)));
            }
        }

        //2. 对Sigma点进行预测（sigmaPre =F*Sigma）
        Matrix Xpre = F.times(sigmaPoint);

        //3.利用2的预测结果计算均值和协方差
        // 均值
        Matrix Xpred_mean = Xpre.subMatrix(Calculation.Ret.NEW, 0, 3, 0, 0).times(weight_0_m);
        for (int i = 1; i < 2 * STATE_DIMENSION + 1; i++) {
            Xpred_mean = Xpred_mean.plus(Calculation.Ret.NEW, true, Xpre.selectColumns(Calculation.Ret.NEW, i)
                    .times(Calculation.Ret.NEW, true, weight_i));
        }
        //协方差
        Matrix Xpred_var = (Xpre.selectColumns(Calculation.Ret.NEW, 0).minus(Calculation.Ret.NEW, true, Xpred_mean))
                .times(Calculation.Ret.NEW, true, Xpre.selectColumns(Calculation.Ret.NEW, 0)
                        .minus(Xpred_mean).transpose()).times(Calculation.Ret.NEW, true, weight_0_c);
        for (int i = 1; i < 2 * STATE_DIMENSION + 1; i++) {
            Xpred_var = Xpred_var.plus((Xpre.selectColumns(Calculation.Ret.NEW, 0).minus(Xpred_mean))
                    .times(Calculation.Ret.NEW, true, Xpre.selectColumns(Calculation.Ret.NEW, 0))
                    .minus(Xpred_mean).transpose()).times(Calculation.Ret.NEW, true, weight_i);
        }
        Xpred_var = Xpred_var.plus(B.times(Q).times(B.transpose()));

        // 4. 根据预测值, 再一次使用UT变换，得到新的sigma点集

        // 5. 观测量预测
        // Z的均值
        Matrix Zpre = Xpre.select(Calculation.Ret.NEW, new long[]{0, 1}, new long[]{0, 2 * STATE_DIMENSION});
//        Zpre.print(Zpre.getRowDimension(), Zpre.getColumnDimension());
        Matrix Zpre_mean = Zpre.select(Calculation.Ret.NEW, new long[]{0, 1}, new long[]{0, 0}).times(Calculation.Ret.NEW, true, weight_0_m);
        for (int i = 1; i < 2 * STATE_DIMENSION + 1; i++) {
            Zpre_mean = Zpre_mean.plus(Zpre.select(Calculation.Ret.NEW, new long[]{0, 1}, new long[]{i, i}).times(Calculation.Ret.NEW, true, weight_i));
        }

        // Z的方差
        Matrix Zpre_var = (Zpre.select(Calculation.Ret.NEW, new long[]{0, 1}, new long[]{0, 0}).minus(Calculation.Ret.NEW, true, Zpre_mean))
                .times(Calculation.Ret.NEW, true, Zpre.select(Calculation.Ret.NEW, new long[]{0, 1}, new long[]{0, 0})
                        .minus(Calculation.Ret.NEW, true, Zpre_mean).transpose()).times(Calculation.Ret.NEW, true, weight_0_c);
        for (int i = 1; i < 2 * STATE_DIMENSION + 1; i++) {
            Zpre_var = Zpre_var.plus((Zpre.select(Calculation.Ret.NEW, new long[]{0, 1}, new long[]{i, i})
                    .minus(Zpre_mean)).times(Calculation.Ret.NEW, true, Zpre.select(Calculation.Ret.NEW, new long[]{0, 1}, new long[]{i, i})
                    .minus(Zpre_mean).transpose()).times(Calculation.Ret.NEW, true, weight_i));
        }
        Zpre_var = Zpre_var.plus(R);

        // 状态与测量间的协方差
        Matrix pxz = (Xpre.selectColumns(Calculation.Ret.NEW, 0).minus(Xpred_mean))
                .times(Calculation.Ret.NEW, true, Zpre.select(Calculation.Ret.NEW, new long[]{0, 1}, new long[]{0, 0})
                        .minus(Zpre_mean).transpose()).times(Calculation.Ret.NEW, true, weight_0_c);
        for (int i = 1; i < 2 * STATE_DIMENSION + 1; i++) {
            pxz = pxz.plus((Xpre.selectColumns(Calculation.Ret.NEW, i).minus(Xpred_mean))
                    .times(Calculation.Ret.NEW, true, Zpre.select(Calculation.Ret.NEW, new long[]{0, 1}, new long[]{i, i})
                            .minus(Zpre_mean).transpose()).times(Calculation.Ret.NEW, true, weight_i));
        }

        Matrix measure = Matrix.Factory.importFromArray(new double[][]{{measurePosition.getX()}, {measurePosition.getY()}});
        K = pxz.times(Calculation.Ret.NEW, true, Zpre_var.inv());

        Matrix UKFresult = Xpred_mean.plus(K.times(measure.minus(Zpre_mean)));
        P = Xpred_var.minus(K.times(Zpre_var).times(K.transpose()));
        measurePosition.setX(UKFresult.getAsDouble(0, 0));
        measurePosition.setY(UKFresult.getAsDouble(1, 0));
        return measurePosition;
    }

    /**
     * @param matrix    原矩阵
     * @param column    被替换的列
     * @param subMatrix 需要替换的列
     * @return
     */
    private void setColumn(Matrix matrix, int column, Matrix subMatrix) {
        for (int i = 0; i < matrix.getRowCount(); i++) {
            matrix.setAsDouble(subMatrix.getAsDouble(i, 0), i, column);
        }
    }
}
